use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        State,
    },
    response::IntoResponse,
};
use futures_util::{sink::SinkExt, stream::StreamExt};
use std::sync::Arc;

use super::state::AppState;

pub async fn ws_handler( ws: WebSocketUpgrade, State(state): State<Arc<AppState>> ) -> impl IntoResponse {
    ws.on_upgrade(|socket| websocket(socket, state))
}
// 此函数处理单个 websocket 连接，即单个连接的客户端/用户，我们将为其生成两个独立的任务（接收/发送聊天消息)
async fn websocket(stream: WebSocket, state: Arc<AppState>) {
    // 通过拆分，我们可以同时发送和接收
    let (mut sender, mut receiver) = stream.split();
    // 循环直到找到一条文本消息
    let mut username = String::new();
    // 循环直到找到一条文本消息
    while let Some(Ok(message)) = receiver.next().await {
        if let Message::Text(name) = message {
            // 如果未取客户端发送的 username，则填写 username 字符串
            check_username(&state, &mut username, &name);
            // 如果不为空，我们想退出循环，否则我们想退出函数
            if !username.is_empty() {
                break;
            } else {
                // 向我们的客户发送用户名已被占用的信息
                let _ = sender
                    .send(Message::Text(String::from("用户名已被占用").into()))
                    .await;
                return;
            }
        }
    }

    // 我们在发送“已加入”消息之前订阅，这样我们就可以显示给我们的客户端
    let mut rx = state.tx.subscribe();

    // 现在向所有订阅者发送“已加入”消息
    let _ = state.tx.send(format!("{username} 进入聊天室"));

    // 生成第一个接收广播消息和发送文本的任务，通过 websocket 向客户端发送消息
    let mut send_task = tokio::spawn(async move {
        while let Ok(msg) = rx.recv().await {
            if sender.send(Message::Text(msg.into())).await.is_err() {
                break;
            }
        }
    });

    // 克隆我们想要传递(移动)给接收任务的东西
    let tx = state.tx.clone();
    let name = username.clone();

    // 生成一个从websocket获取消息的任务，前置用户名称，并发送给所有广播订阅者
    let mut recv_task = tokio::spawn(async move {
        while let Some(Ok(Message::Text(text))) = receiver.next().await {
            // Add username before message.
            let _ = tx.send(format!("{name}说: {text}"));
        }
    });

    // 如果其中任何一个任务运行到完成，我们就中止另一个任务
    tokio::select! {
        _ = &mut send_task => recv_task.abort(),
        _ = &mut recv_task => send_task.abort(),
    };

    // 用户离开消息
    let _ = state.tx.send(format!("{username} left."));
    // 从 Map 中删除用户名，以便新客户端可以再次使用它
    state.user_set.lock().unwrap().remove(&username);
}

fn check_username(state: &AppState, string: &mut String, name: &str) {
    let mut user_set = state.user_set.lock().unwrap();

    if !user_set.contains(name) {
        user_set.insert(name.to_owned());
        string.push_str(name);
    }
}